# encoding: utf-8
"""
@Author:      "June"
@CreateTime:  2020/5/28 21:17 
@File:        douyin
@Project:     PyCharm
"""
import json
import os
import random
import time
from imp import reload
from urllib.parse import quote
import frida
import pymongo
import requests
import logging
import config

BASE_DIR = os.path.abspath(os.path.dirname(__file__))
jscode = """
rpc.exports = {
    getsig: function(timestamp, url, device_list, device_id) {
        var sig = "";
        send('enter...');
        Java.perform(
            function() {
                var userInfo = Java.use("com.ss.android.common.applog.UserInfo");
                var StcSDKFactory = Java.use("com.ss.sys.ces.out.StcSDKFactory");
                var GlobalContext = Java.use("com.ss.android.common.applog.GlobalContext");
                var AwemeApplication = Java.use("com.ss.android.ugc.aweme.app.AwemeApplication");
                var EagleEye = Java.use("com.ss.android.common.applog.EagleEye");
    
                var str = userInfo.getUserInfo(timestamp, url, device_list, device_id);
                var len = str.length;
                var i2 = len / 2;
                var as = str.substring(0, 22);
                var cp = str.substring(22, len);
    
                var context = GlobalContext.getContext();
                var sdk = StcSDKFactory.getSDK(context, AwemeApplication.p().m());
    
                var mas = EagleEye.byteArrayToHexStr(sdk.encode(getBytes(as)));
                
                //send(as);
                //send(cp);
                //send(mas);
                sig = '&as='+as+'&cp='+cp+'&mas='+mas;
            }
        )
        return sig;
    }
}

function getBytes(s) {
    var bytes = [];
        for (var i = 0; i < s.length; i++) {
            bytes.push(s.charCodeAt(i));
        }
    return bytes;
}
"""


class DouYin:
    def __init__(self, keyword):
        self.headers = {
            'User-Agent': "com.ss.android.ugc.aweme.lite/203 (Linux; U; Android 5.0.1; zh_CN; MX4 Pro; Build/LRX22C; Cronet/58.0.2991.0)",
            'X-Gorgon': '03006cc04005ea0986339c621ab0abc35940f5e2c14d8a108012',
            'X-Khronos': '1590732699',
            # 'X-SS-REQ-TICKET': '1590588862616',
            # 'X-SS-TC': '0',
        }

        self.connect = pymongo.MongoClient(host=config.host,
                                           port=config.port)
        self.db = self.connect['douyin']
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(logging.DEBUG)
        if not isinstance(keyword, list):
            raise Exception('Parameter keyword error,"list" required.')
        self.keyword = keyword
        self.device_id = None
        self.device_id = None
        self.sig = None
        self.ts = None
        self.rticket = None

    def on_message(self, message, data):
        if message['type'] == 'send':
            print("[*] {0}".format(message['payload']))
        else:
            print(message)

    def get_script(self):
        process = frida.get_usb_device().attach('com.ss.android.ugc.aweme.lite')
        script = process.create_script(jscode)
        script.on('message', self.on_message)
        script.load()
        return script

    def get_sig(self, script):
        """
        获取请求参数，as,cp,mas
        从getUserInfo文件中的hook函数可以看出需要4个参数，
        分别是时间戳timestamp，请求URL，设备信息列表device_list，设备号device_id
        :return: as=xxx&ps=xxx&mas=xxx
        """
        timestamp_w = time.time()
        rticket = int(timestamp_w * 1000)
        ts = int(timestamp_w)
        self.device_id = ''.join(random.sample(['1', '2', '3', '4', '5', '6', '7', '8', '9', '3', '5', '1'], 11))
        url = f'https://aweme.snssdk.com/aweme/v1/challenge/fresh/aweme/?ch_id=1665117510154247&' \
            f'query_type=0&cursor=0&count=20&type=5&retry_type=no_retry&iid=3562061066669069&' \
            f'device_id=69133478649&ac=wifi&channel=xiaomi&aid=2329&app_name=douyin_lite&' \
            f'version_code=203&version_name=2.0.3&device_platform=android&ssmix=a&' \
            f'device_type=MI+5s&device_brand=Xiaomi&language=zh&os_api=23&os_version=6.0.1&' \
            f'uuid=490000000185949&openudid=3cffa76ab4f67a77&manifest_version_code=203&' \
            f'resolution=810*1440&dpi=270&update_version_code=2030&_rticket={rticket}&ts={ts}'

        device_list = ['os_api', '23', 'device_type', 'MI 5s', 'device_platform',
                       'android', 'ssmix', 'a', 'iid', '3562061066669069', 'manifest_version_code',
                       '203', 'dpi', '270', 'uuid', '490000000185949', 'version_code', '203',
                       'app_name', 'douyin_lite', 'version_name', '2.0.3', 'openudid',
                       '3cffa76ab4f67a77', 'device_id', self.device_id, 'resolution', '810*1440',
                       'os_version', '6.0.1', 'language', 'zh', 'device_brand', 'Xiaomi',
                       'ac', 'wifi', 'update_version_code', '2030', 'aid', '2329', 'channel', 'xiaomi']
        sig = script.exports.getsig(ts, url, device_list, self.device_id)
        self.sig, self.ts, self.rticket = sig, ts, rticket
        # return sig, ts, rticket

    def get_category_list(self, sig, ts, rticket):
        for kwd in self.keyword:
            kwd = quote(kwd)
            # 关键字搜索URL，可获取user_id
            # keyword_url = 'https://aweme.snssdk.com/aweme/v1/discover/search/?' \
            #                   f'cursor=0&keyword={kwd}&count=10&type=1&' \
            #                   f'ts={ts}&app_type=lite&os_api=23&' \
            #               'device_type=MI%205s&device_platform=android&' \
            #               'ssmix=a&iid=3562061066669069&manifest_version_code=203&' \
            #               'dpi=270&uuid=490000000185949&version_code=203&' \
            #               'app_name=douyin_lite&version_name=2.0.3&' \
            #               f'openudid=3cffa76ab4f67a77&device_id={self.device_id}&' \
            #               'resolution=810*1440&os_version=6.0.1&language=zh&' \
            #               'device_brand=Xiaomi&ac=wifi&update_version_code=2030&' \
            #                   f'aid=2329&channel=xiaomi&_rticket={rticket}' + sig
            start_time = time.time()
            while True:
                if config.close:
                    self.logger.warn('手动停止.')
                    break
                url = f'https://api.amemv.com/aweme/v1/category/list/?' \
                    f'cursor=0&count=5&retry_type=no_retry&iid=3562061066669069&device_id={self.device_id}&' \
                    f'ac=wifi&channel=xiaomi&aid=2329&app_name=douyin_lite&version_code=203&' \
                    f'version_name=2.0.3&device_platform=android&ssmix=a&device_type=MI+5s&' \
                    f'device_brand=Xiaomi&language=zh&os_api=23&os_version=6.0.1&' \
                    f'uuid=490000000185949&openudid=3cffa76ab4f67a77&manifest_version_code=203&' \
                    f'resolution=810*1440&dpi=270&update_version_code=2030&_rticket={rticket}&' + sig
                result = requests.get(url, headers=self.headers).json()
                time.sleep(1)
                category_list = result.get('category_list')
                if category_list:
                    for item in category_list:
                        self.logger.warn('获取数据成功====>'+json.dumps(item))
                        yield item
                    # 每隔5秒刷新一次
                    time.sleep(5)
                else:
                    self.logger.warn('获取数据失败====>'+json.dumps(result))
                    self.logger.warn(url)
                    return False
                # 每隔30秒重载下config文件
                if time.time() - start_time >= 30:
                    reload(config)
                    start_time = time.time()

    def get_detail(self, category_list):
        # url = 'https://aweme.snssdk.com/aweme/v1/aweme/post/?' \
        #           f'user_id={uid}&max_cursor=0&count=20' \
        #       '&retry_type=no_retry&iid=3562061066669069' \
        #       f'&device_id={self.device_id}&ac=wifi&channel=xiaomi' \
        #       '&aid=2329&app_name=douyin_lite&version_code=203' \
        #       '&version_name=2.0.3&device_platform=android' \
        #       '&ssmix=a&device_type=MI+5s&device_brand=Xiaomi' \
        #       '&language=zh&os_api=23&os_version=6.0.1' \
        #       '&uuid=490000000185949&openudid=3cffa76ab4f67a77' \
        #       '&manifest_version_code=203&resolution=810*1440&dpi=270' \
        #           f'&update_version_code=2030&_rticket={rticket}' \
        #           f'&ts={ts}' + sig
        # result = requests.get(url, headers=self.headers).json()
        #     time.sleep(1)
        #     aweme_list = result.get('aweme_list')
        #     if aweme_list:
        #         for item in aweme_list:
        #             self.save_to_mongodb(item)
        #             self.logger.warn('保存详情数据成功====>'+json.dumps(item))
        #             return True
        #     self.handle()
        # self.logger.warn('获取详情数据失败====>'+json.dumps(result))
        # return False
        for item in category_list:
            self.save_to_mongodb(item)

    def save_to_mongodb(self, item):
        self.db['douyin'].update({'_id': item['challenge_info']['cid']}, {'$set': item}, True)
        self.logger.warn('数据保存成功===>'+json.dumps(item))

    def handle(self):
        script = self.get_script()
        self.get_sig(script)

    def main(self):
        self.handle()
        self.logger.warn('sig='+self.sig)
        self.logger.warn('ts='+str(self.ts))
        self.logger.warn('rticket='+str(self.rticket))
        self.logger.warn('device_id='+self.device_id)
        category_list = self.get_category_list(self.sig, self.ts, self.rticket)
        if category_list:
            self.get_detail(category_list)
        else:
            self.main()



if __name__ == '__main__':
    keyword = ['妹子']
    douyin = DouYin(keyword)
    douyin.main()

